from testCreateSM2keys import *
from gmssl import sm2, sm3, func
import binascii
import os


def sm3_hash_and_sign_dict(data_dict, private_key_hex, public_key_hex):
    """
    对输入的字典数据进行SM3哈希计算，并使用SM2算法进行签名，返回哈希结果（十六进制字符串）和签名结果（十六进制字符串）。
    需要调用者传入十六进制格式的私钥和公钥。

    参数:
    data_dict (dict): 需要处理的字典数据
    private_key_hex (str): 十六进制格式的私钥
    public_key_hex (str): 十六进制格式的公钥

    返回:
    tuple: 包含哈希结果（十六进制字符串）和签名结果（十六进制字符串）的元组
    """
    # 将字典数据转换为字节类型
    data_bytes = str(data_dict).encode('utf-8')
    # 计算SM3哈希值并转换为十六进制字符串格式
    hash_value = sm3.sm3_hash(func.bytes_to_list(data_bytes))

    # 校验传入的私钥和公钥是否为十六进制字符串且长度符合基本要求（简单校验，实际可更严格）
    if not all(c in "0123456789abcdefABCDEF" for c in private_key_hex):
        raise ValueError("私钥不是合法的十六进制字符串格式")
    if not all(c in "0123456789abcdefABCDEF" for c in public_key_hex):
        raise ValueError("公钥不是合法的十六进制字符串格式")

    # 创建SM2对象实例，传入公私钥
    sm2_crypt = sm2.CryptSM2(public_key=public_key_hex, private_key=private_key_hex)
    # 进行签名操作，获取十六进制格式的签名结果
    sign_result = sm2_crypt.sign_with_sm3(data_bytes)
    return hash_value, sign_result


def verify_dict_signature(data_dict, sign_value_hex, public_key_hex):
    """
    验证给定字典数据对应的签名是否有效。
    需要调用者传入十六进制格式的私钥、公钥以及签名值。

    参数:
    data_dict (dict): 原始的字典数据
    sign_value_hex (str): 需要验证的签名值（十六进制字符串形式）
    public_key_hex (str): 十六进制格式的公钥

    返回:
    bool: 验证结果，True表示签名有效，False表示签名无效
    """
    # 将字典数据转换为字节类型
    data_bytes = str(data_dict).encode('utf-8')

    # 校验传入的公钥和签名值是否符合格式要求
    if not all(c in "0123456789abcdefABCDEF" for c in public_key_hex):
        raise ValueError("公钥不是合法的十六进制字符串格式")
    if not all(c in "0123456789abcdefABCDEF" for c in sign_value_hex):
        raise ValueError("签名值不是合法的十六进制字符串格式")

    # 创建SM2对象实例，传入公钥和占位私钥
    sm2_crypt = sm2.CryptSM2(public_key=public_key_hex, private_key=" ")
    # 进行验证操作
    verify_result = sm2_crypt.verify_with_sm3(sign_value_hex, data_bytes)
    return verify_result


def encrypt_sm4_key_with_public_key(sm4_key, public_key_hex):
    """
    使用给定的十六进制格式的公钥对SM4密钥进行加密，并返回字节类型的加密结果。

    参数:
    sm4_key (bytes): 需要加密的SM4密钥，字节类型
    public_key_hex (str): 十六进制格式的公钥

    返回:
    bytes: 加密后的SM4密钥数据
    """
    # 校验传入的公钥是否为十六进制字符串且长度符合基本要求（简单校验，实际可更严格）
    if not all(c in "0123456789abcdefABCDEF" for c in public_key_hex):
        raise ValueError("公钥不是合法的十六进制字符串格式")

    # 创建SM2对象实例，传入公钥（私钥占位）
    sm2_crypt = sm2.CryptSM2(public_key=public_key_hex, private_key='')
    # 进行加密操作
    encrypted_sm4_key = sm2_crypt.encrypt(sm4_key)
    return encrypted_sm4_key


def decrypt_sm4_key_with_private_key(encrypted_sm4_key_bytes, private_key_hex):
    """
    使用给定的十六进制格式的私钥对字节类型的加密后的SM4密钥进行解密，并返回十六进制字符串格式的解密结果。

    参数:
    encrypted_sm4_key_bytes (bytes): 字节类型的加密后的SM4密钥数据
    private_key_hex (str): 十六进制格式的私钥

    返回:
    str: 十六进制格式的解密后的SM4密钥
    """
    # 校验传入的私钥是否为十六进制字符串且长度符合基本要求（简单校验，实际可更严格）
    if not all(c in "0123456789abcdefABCDEF" for c in private_key_hex):
        raise ValueError("私钥不是合法的十六进制字符串格式")

    # 创建SM2对象实例，传入私钥和占位公钥
    sm2_crypt = sm2.CryptSM2(public_key='', private_key=private_key_hex)
    # 进行解密操作
    decrypted_sm4_key = sm2_crypt.decrypt(encrypted_sm4_key_bytes)
    return binascii.hexlify(decrypted_sm4_key).decode('utf-8')


if __name__ == "__main__":
    # 示例字典数据
    test_dict = {"key1": "value1", "key2": "value2"}
    # 示例十六进制格式的私钥和公钥（假设已有生成公私钥对的函数create_key_pair，实际需替换为真实有效的方式获取）
    private_key_hex_example, public_key_hex_example = create_key_pair()

    # 1. 进行SM3哈希和SM2签名操作
    hash_result, sign_result = sm3_hash_and_sign_dict(test_dict, private_key_hex_example, public_key_hex_example)
    print("SM3哈希结果:", hash_result)
    print("SM2签名结果:", sign_result)

    # 2. 进行签名验证操作
    is_valid_signature = verify_dict_signature(test_dict, sign_result, public_key_hex_example)
    print("签名验证结果:", is_valid_signature)

    # 3. 示例SM4密钥生成（16字节随机字节数据符合SM4密钥长度要求）
    sm4_key_example = os.urandom(16)

    # 4. 使用公钥加密SM4密钥
    encrypted_sm4_key_result = encrypt_sm4_key_with_public_key(sm4_key_example, public_key_hex_example)
    print("加密后的SM4密钥（字节类型）:", encrypted_sm4_key_result)

    # 5. 使用私钥解密SM4密钥
    decrypted_sm4_key_result = decrypt_sm4_key_with_private_key(encrypted_sm4_key_result, private_key_hex_example)
    print("解密后的SM4密钥（十六进制格式）:", decrypted_sm4_key_result)